package jsmile.utils;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import smile.Network;
import smile.learning.DataMatch;
import smile.learning.DataSet;
import smile.learning.EM;

public class Learn {

	protected static String getDBNVarName(String plainVariable) {
		return plainVariable.split("_")[0];
	}

	/**
	 * Default dataset loading method in SMILE doesn't treat data the same way
	 * as GENIE does. It doesn't recognize variables O, O_1, O_2 ... as a
	 * sequence of O. Instead it treats them as independent variables. 
	 * 
	 * @param net
	 * @param smileDataFile
	 * @param range range of example rows from the source dataset to use for learning 
	 * @return
	 */
	public static DataSet loadAndFixDBNDataset(Network net, String smileDataFile, int[] range) {
		DataSet ds = new DataSet();
		ds.readFile(smileDataFile);
		
		// join lists of outcomes
		Map<String, Set<String>> variableToOutcomes = new HashMap<>();
		for(int i = 0; i < ds.getVariableCount(); i++) {
			String var = getDBNVarName(ds.getVariableId(i));
			Set<String> outcomes = variableToOutcomes.get(var);
			if(outcomes == null) {
				outcomes = new HashSet<>();
				variableToOutcomes.put(var, outcomes);
			}
			outcomes.addAll(Arrays.asList(ds.getStateNames(i)));
		}
		
		// create new data set
		DataSet data = new DataSet();
				
		// sort outcomes list and convert it into array
		Map<String, String[]> variableToOutcomesArr = new HashMap<>();
		for(Entry<String, Set<String>> entry : variableToOutcomes.entrySet()) {
			String nodeId = entry.getKey();
			String[] outcomes = new ArrayList<>(entry.getValue()).toArray(new String[0]);
			Arrays.sort(outcomes);
			variableToOutcomesArr.put(nodeId, outcomes);
		
			// check that domains in dataset and network match
			String[] expectedOutcomes = net.getOutcomeIds(nodeId);
			if(!Arrays.equals(outcomes, expectedOutcomes)) {
				throw new IllegalArgumentException("Domain of " + nodeId + " does not match the data.");
			}
		}
			
		// fill data in new fixed dataset
		// init variables
		for(int j = 0; j < ds.getVariableCount(); j++) {
			//int varId = j % match.length;
			String varStr = ds.getVariableId(j);
			data.addIntVariable(varStr);
			data.setStateNames(j, variableToOutcomesArr.get(getDBNVarName(varStr)));
		}
		
		// set the data
		for (int i = range[0]; (i < ds.getRecordCount()-range[0]) && (i < range[1]); i++) {
			data.addEmptyRecord();
			for(int j = 0; j < ds.getVariableCount(); j++) {
				int valueOrig = ds.getInt(j, i);
				String valueStr = ds.getStateNames(j)[valueOrig];
				// find matching value in fixed
				int valueFix = Arrays.binarySearch(data.getStateNames(j), valueStr);
				
				data.setInt(j, i-range[0], valueFix);
			}
		}
		return data;
	}
	/**
	 * Default creation of datamatch in smile handles only variables in the first time slice.
	 * This method adds matches for other time slices.
	 * @param data
	 * @param net
	 * @return
	 */
	public static DataMatch[] createFixedDBNDataMatch(DataSet data, Network net) {
		DataMatch[] match = data.matchNetwork(net); //-> do i need this line? what for? It appears me as void, not creating a DataMatch[]...
		
		// fix data match for DBNs
		DataMatch[] finalMatch = new DataMatch[data.getVariableCount()];
		for(int i = 0; i < finalMatch.length; i++) {
			int node = i % match.length;
			DataMatch orig = match[node];
			finalMatch[i] = new  DataMatch(i, orig.node, i / match.length);
		}
		return finalMatch;
	}

	public static void learnEM(Network net, String smileDataFile) {
		learnEM(net, smileDataFile, new int[] {0, Integer.MAX_VALUE});
	}
	
	public static void learnEM(Network net, String smileDataFile, int[] range) {
		DataSet ds = loadAndFixDBNDataset(net, smileDataFile, range); 
		//ds.readFile(smileDataFile);
		
		DataMatch[] match = createFixedDBNDataMatch(ds, net);

		final EM em = new EM();
		em.setEqSampleSize(0); //-> learning algorithm parameter, if i'm right
		
		em.setSeed(1);
		//em.setRandomizeParameters(true); //-> Not sure about this, it means that randomize the values of the nodes that are not informed in the txt?
		
		em.setUniformizeParameters(true);
		em.setRelevance(true);
		
		long start = System.nanoTime();
		em.learn(ds, net, match);
		System.out.println("Learning took: " + (System.nanoTime() - start) + " nanoseconds.");
	}
}
